#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2025-02-01-0:20
# @Author  : 
# @File    : test-xml.py
# @Software: PyCharm

import os
import json

'''人为构造xml文件的格式'''
out0 = '''<annotation>
    <filename>%(name)s</filename>
    <source>
        <database>None</database>
    </source>
    <size>
        <width>%(width)d</width>
        <height>%(height)d</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>
'''
out1 = '''    <object>
        <name>%(class)s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>%(xmin)d</xmin>
            <ymin>%(ymin)d</ymin>
            <xmax>%(xmax)d</xmax>
            <ymax>%(ymax)d</ymax>
        </bndbox>
    </object>
'''

out2 = '''</annotation>
'''


def annos2xml(annos, out_dir):
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    source = {}  # 用于暂存图片名字、尺寸等基本信息
    h, w = 2048, 2048  # 图片大小
    image_name = annos["path"]  # eg test/10056.jpg
    out_xml = os.path.join(out_dir, image_name[:-4] + '.xml')
    fxml = open(out_xml, 'w')  # 新建对应 xml 文件
    source['name'] = image_name  # 图片路径

    source['width'] = w  # 图片宽度
    source['height'] = h  # 图片高度

    fxml.write(out0 % source)  # 写入 out0 模板

    label = {}  # 暂存当前图片的 box 信息
    for obj in annos['objects']:  # 遍历每个子列表即原 txt 的每一行
        label['class'] = obj['category']  # 写入编号对应类别标签

        '''读取框的坐标'''
        label['xmin'] = obj['bbox']['xmin']
        label['ymin'] = obj['bbox']['ymin']
        label['xmax'] = obj['bbox']['xmax']
        label['ymax'] = obj['bbox']['ymax']

        fxml.write(out1 % label)  # 写入 out1 模板
    fxml.write(out2)


def main():
    filedir = "./annotations_all.json"  # json文件
    ids_file = "./test/ids.txt"  # 指定为 train/test
    ids = open(ids_file).read().splitlines()  # 获取 id 编号
    annos = json.loads(open(filedir).read())
    outdir = './'  # xml 保存目录
    img_num = len(ids)  # 统计当前图片数量
    cnt_id = 1  # 计数变量
    for imgid in ids:
        print('\rprocessing :[{} / {}]'.format(cnt_id, img_num), end="")
        cnt_id += 1

        # 跳过没有目标的图片
        if imgid in annos['imgs']:
            xml_cls = annos['imgs'][imgid]
            annos2xml(xml_cls, outdir)  # 生成 xml 文件
    ''' 
    annos['imgs'][imgid] 示例：
    {'path': 'test/10056.jpg', 
    'id': 10056, 
    'objects': [
    {'bbox': {'xmin': 452.475, 'ymin': 886.139, 'xmax': 468.3168, 'ymax': 916.8317}, 'category': 'i5'}, 
    {'bbox': {'xmin': 1274.26, 'ymin': 927.723, 'xmax': 1294.0594, 'ymax': 949.505}, 'category': 'i5'}, 
    {'bbox': {'xmin': 414.851, 'ymin': 877.228, 'xmax': 431.6832, 'ymax': 909.901}, 'category': 'pne'}, 
    {'bbox': {'xmin': 1215.84, 'ymin': 928.713, 'xmax': 1237.6237999999998, 'ymax': 950.495}, 'category': 'pne'}, 
    {'bbox': {'xmin': 2016.83, 'ymin': 910.891, 'xmax': 2032.6733, 'ymax': 934.6535}, 'category': 'pne'}]}
    '''


if __name__ == "__main__":
    main()
